{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization Dynamics: an Interactive Tool\n",
    "\n",
    "This notebook concentrates on improving your intuitive understanding of the methods employed by\n",
    "[Cooper](https://github.com/cooper-org/cooper) for solving constrained optimization problems.\n",
    "\n",
    "To achieve so, we consider a [toy constrained optimization problem](#cmp) in 2D.\n",
    "We provide an interactive widget which shows the optimization path realized when solving the problem using [Cooper](https://github.com/cooper-org/cooper), for a variety of hyper-parameter settings.\n",
    "We will guide you through specific settings which highlight some interesting properties of gradient descent-ascent schemes for solving min-max Lagrangian based problems.\n",
    "You can also play with the widget to explore the behavior of optimization on your own.\n",
    "\n",
    "If you are interested in a tutorial on how to use [Cooper](https://github.com/cooper-org/cooper),\n",
    "visit [this tutorial](https://github.com/cooper-org/cooper/tree/master/tutorials/logistic_regression.ipynb).\n",
    "For a fully realized application in the context of deep learning, visit\n",
    "[this repository](https://github.com/gallego-posada/constrained_l0).\n",
    "\n",
    "> **Acknowledgement**\n",
    "> The presented visualizations and optimization problems follow\n",
    "> closely the blogposts by [Degrave and Korshunova (2021a, 2021b)](#blog1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TODO: update TOC\n",
    "## Table of Contents:\n",
    "* [Setup](#setup)\n",
    "* [Constrained Minimization Problem](#cmp)\n",
    "* [Widget](#widget)\n",
    "* [References](#references)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup <a class=\"anchor\" id=\"setup\"></a>\n",
    "Install Cooper, with the requirements for running `examples`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -e git+https://github.com/cooper-org/cooper#egg=.[examples]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "import cooper\n",
    "\n",
    "from widget import Toy2DWidget\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "torch.manual_seed(0)\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constrained Minimization Problem <a class=\"anchor\" id=\"cmp\"></a>\n",
    "\n",
    "Consider the following constrained optimization problem on\n",
    " the 2D domain $(x, y) \\in [0,\\pi/2] \\times [0,\\infty]$\n",
    "\n",
    "$$\\begin{align*}\n",
    "\\underset{x, y}{\\text{min}}\\quad f(x,y) &:= \\left(1 - \\text{sin}(x) \\right) \\ \\big(1+(y - 1)^2\\big) & \\tag{1} \\\\\n",
    "s.t. \\quad  g(x,y) &:= \\left(1 - \\text{cos}(x) \\right)\\ \\big(1+(y-1)^2\\big) - \\epsilon \\leq 0 & \\\\\n",
    "\\end{align*}$$\n",
    "\n",
    "given some $\\epsilon \\geq 0 $.\n",
    "Note how both $f$ and $g$ are convex functions in the specified domain.\n",
    "As such, this constrained minimization problem is a convex problem.\n",
    "\n",
    "The following class implements this CMP:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Toy2DCMP(cooper.ConstrainedMinimizationProblem):\n",
    "    def __init__(self, is_constrained=False, epsilon=1.0, problem_type=\"Convex\"):\n",
    "        self.problem_type = problem_type\n",
    "        self.epsilon = epsilon\n",
    "        super().__init__(is_constrained)\n",
    "\n",
    "    def closure(self, params):\n",
    "        \"\"\"This function evaluates the objective function and constraint\n",
    "        defect. It updates the attributes of this CMP based on the results.\"\"\"\n",
    "\n",
    "        x, y = params[:, 0], params[:, 1]\n",
    "\n",
    "        if self.problem_type == \"Convex\":\n",
    "            f = (1 - torch.sin(x)) * (1 + (y - 1.0) ** 2)\n",
    "            # In standard form (defect <= 0)\n",
    "            g = (1 - torch.cos(x)) * (1 + (y - 1.0) ** 2) - self.epsilon\n",
    "        elif self.problem_type == \"Concave\":\n",
    "            f = torch.sin(x) * (1 + (y - 1.0) ** 2)\n",
    "            # in standard form (defect <= 0)\n",
    "            g = torch.cos(x) * (1 + (y - 1.0) ** 2) - self.epsilon\n",
    "        else:\n",
    "            raise ValueError(\"Unknown problem type.\")\n",
    "\n",
    "        # Store the values in a CMPState as attributes\n",
    "        state = cooper.CMPState(loss=f, ineq_defect=g)\n",
    "\n",
    "        return state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that [Cooper](https://github.com/cooper-org/cooper) employs a flexible approach to solving problems like Eq. (1). First, it formulates the Lagrangian associated with the optimization problem. Then, it employs saddle-point optimizing techniques on the Lagrangian's min-max game. The Lagrangian and its respective game are as follows:\n",
    "\n",
    "$$ \\begin{align*}\n",
    "& \\underset{x, y}{\\text{min}}\\ \\underset{\\lambda \\geq 0}{\\text{max}}\\ f(x,y) +\n",
    "\\lambda \\, g(x,y) \\tag{2} \\\\\n",
    "& \\underset{x, y}{\\text{min}}\\ \\underset{\\lambda \\geq 0}{\\text{max}}\\,\n",
    "\\big(1 - \\text{sin}(x) \\big) \\ \\big(1+(y-1)^2\\big) +\n",
    "\\lambda \\big(1 - \\text{cos}(x) \\big)\\ \\big(1+(y-1)^2\\big)  \\\\\n",
    "\\end{align*} $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bonus: non-convex $f(x, y)$ and $g(x, y)$.\n",
    "\n",
    "Associated with `Toy2DCMP(problem_type=\"Concave\")`\n",
    "\n",
    "Consider a similar optimization problem to Eq. (1), also on $[0,\\pi/2] \\times [0,\\infty]$:\n",
    "\n",
    "$$\\begin{align*}\n",
    "\\underset{x, y}{\\text{min}}\\quad f(x,y) &:= \\text{sin}(x) \\ \\big(1+(y - 1)^2\\big) & \\tag{3} \\\\\n",
    "s.t. \\quad  g(x,y) &:= \\text{cos}(x)\\ \\big(1+(y-1)^2\\big) - \\epsilon \\leq 0 & \\\\\n",
    "\\end{align*}$$\n",
    "\n",
    "given some $\\epsilon \\geq 0 $.\n",
    "$f$ and $g$ are concave functions with respect to $x$ in this case."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Widget <a class=\"anchor\" id=\"widget\"></a>\n",
    "\n",
    "The `Toy2DWidget` class implements a widget which shows the optimization dynamics of [Cooper](https://github.com/cooper-org/cooper) when solving the problems in Eq. (1) and Bonus: Eq. (3). The widget keeps track of various metrics throughout the optimization of Eq. (2) via simultaneous gradient descent on the parameters $(x,y)$ and gradient ascent on the dual variable $\\lambda$.\n",
    "\n",
    "You can play by adjusting different hyper-parameters (optimizer type, learning rate) of the widget to see how they affect the optimization path.\n",
    "\n",
    "This first widget is flexible as it allows you to tune all of the main components required for doing optimization with [Cooper](https://github.com/cooper-org/cooper).\n",
    "In the following sections of this notebook, we will guide your exploration to help develop some intuitions with respect to the optimization methods behind [Cooper](https://github.com/cooper-org/cooper). Feel free to come back here at any moment to try out interesting dynamics!\n",
    "\n",
    "Run the widget to get a sample trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "da954264ccbd44159c0a97fbb8d9eb32",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HBox(children=(Dropdown(description='Problem type', options=('Convex', 'Concave'), value='Conve…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "widget = Toy2DWidget(Toy2DCMP)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reading the plots\n",
    "\n",
    "Run the next cell. This widget has some hyper-parameters fixed, but you can still tune the constraint level $\\epsilon$ and the initialization.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cd0d7b822b2346e2880c1d36f5483a9e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HBox(children=(FloatSlider(value=0.7, description='Const. level', max=1.5, min=-0.2, step=0.05)…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "widget = Toy2DWidget(\n",
    "    Toy2DCMP,\n",
    "    num_iters=300,\n",
    "    problem_type=\"Convex\",\n",
    "    primal_optim=\"Adam\",\n",
    "    primal_lr=0.02,\n",
    "    dual_optim=\"SGD\",\n",
    "    dual_lr=0.05,\n",
    "    x=0.7,\n",
    "    y=2.0,\n",
    "    extrapolation=False,\n",
    "    dual_restarts=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**First column**\\\n",
    "On the top left, the objective $f$ is displayed. In this example, it decreases along training, but not on a monotonic fashion.\n",
    "The value of the constraint $g$ and the Lagrange multiplier $\\lambda$ are presented on the bottom left panel. The constraint is asymptotically satisfied for $\\epsilon = 0.7$.\n",
    "\n",
    "**Middle column**\\\n",
    "The middle plot shows the optimization path followed in $(x, y)$ space. The feasible set is highlighted on blue. Level curves for the loss are presented in gray. For this specific run, parameters constantly come in and out of the feasible region, whilst consistently moving towards the optimum.\n",
    "\n",
    "**Last column**\\\n",
    "The panel on the right shows the optimization path followed in $(f, g)$ space. As such, points to the left of the plot correspond to solutions with low objective values and points towards the bottom correspond to solutions with lower constraint values. The black curve represents the Pareto front of non-dominated solutions on the joint optimization of $(f,g)$. Depending on the constraint level set, the optimal solution of the constrained optimization problem falls in different points of the Pareto front. Finally, the highlighted blue region shows the feasible region of the considered problem.\n",
    "\n",
    "> **Exercise**\n",
    "> * Set $\\epsilon > 1$. What do you observe? Is the multiplier increasing? At what rate does the value of the multiplier change?\n",
    "> * Set $\\epsilon < 0$. At which iteration is the constraint satisfied?\n",
    "> * Move $\\epsilon$ in $(0, 1)$. What is the resulting objective value when the constraint is tighter? How does the multiplier behave?\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Lagrangian multipliers\n",
    "\n",
    "The effective use of [Cooper](https://github.com/cooper-org/cooper) requires an understanding of the behavior of the Lagrangian multipliers $\\lambda$ associated with each constraint. Recall from Eq. (2) the problem we are solving:\n",
    "$$ \\begin{align*}\n",
    "& \\underset{x, y}{\\text{min}}\\ \\underset{\\lambda \\geq 0}{\\text{max}}\\ f(x,y) +\n",
    "\\lambda \\, g(x,y) \\\\\n",
    "\\end{align*} $$\n",
    "\n",
    "\n",
    "A large value of $\\lambda$ is reflected in a large relative contribution of $g$ towards the update of $(x, y)$. Conversely, when $\\lambda = 0$, only gradients of $f$ will be used on a descent step on $(x, y)$.\n",
    "\n",
    "Additionally, the (gradient ascent) update on the multiplier is:\n",
    "$$ \\begin{align*}\n",
    "    \\lambda_{t+1} = \\lambda_t + \\eta_{\\lambda}  \\, g(x_t,y_t) \\\\\n",
    "\\end{align*} $$\n",
    "Therefore, changes on the multiplier depend directly on the value of the constraint, as mediated by the learning rate $\\eta_{\\lambda}$. The value of the learning rate will be reflected in the eagerness to satisfy the constraint. As with other hyper-parameters, choosing its value is an important and complex task.\n",
    "\n",
    "\n",
    "> **Note**\n",
    "> The actual update on $\\lambda$ differs based on the dual optimizer chosen. Nonetheless, the core intuition between the relationship of the constraint value and the multiplier still holds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cf89820ca3ca4c1293245a48c2880e65",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HBox(children=(IntSlider(value=200, description='Max Iters', max=3000, min=100, step=100), Floa…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "widget = Toy2DWidget(\n",
    "    Toy2DCMP,\n",
    "    num_iters=200,\n",
    "    epsilon=0.7,\n",
    "    problem_type=\"Convex\",\n",
    "    primal_optim=\"SGD\",\n",
    "    primal_lr=0.02,\n",
    "    dual_optim=\"SGD\",\n",
    "    x=0.7,\n",
    "    y=2.0,\n",
    "    extrapolation=False,\n",
    "    dual_restarts=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The multiplier does not increase during the first few iterations as the constraint is initially satisfied. During this period, the optimization focuses on the objective function, resulting in a violation of the constraint. The multiplier then kicks in to balance the importance $f$ and $g$. Thereafter, the constraint presents dampened oscillations around 0.\n",
    "\n",
    "\n",
    "> **Exercise**\n",
    "> * Try different values of the dual learning rate. What do you observe? \n",
    "> * Are there oscillations in the value of the constraint or multiplier? Does their amplitude or frequency change for different dual learning rates?\n",
    "> * Is the constraint satisfied by the end of training?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dual Restarts\n",
    "\n",
    "> **TODO**\n",
    "> Dual restarts seems to be detrimental in this specific problem. Maybe because $f$ and $g$ present a nearly orthogonal trade-off?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "561309ae7c4542019781c421fd780bbc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HBox(children=(IntSlider(value=1000, description='Max Iters', max=3000, min=100, step=100), Che…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "widget = Toy2DWidget(\n",
    "    Toy2DCMP,\n",
    "    num_iters=1000,\n",
    "    epsilon=0.7,\n",
    "    problem_type=\"Convex\",\n",
    "    primal_optim=\"SGD\",\n",
    "    primal_lr=0.02,\n",
    "    dual_optim=\"SGD\",\n",
    "    dual_lr=0.5,\n",
    "    x=0.7,\n",
    "    y=1.8,\n",
    "    extrapolation=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extrapolation\n",
    "\n",
    "The are no convergence guarantees for gradient descent-ascent updates on general min-max games. Korpelevich (1976) and Gidel et. al (2019)\n",
    "\n",
    "> **TODO**\n",
    "> Non-convex case often gets stuck at the corners. Otherwise, it oscillates even for Extrapolation!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02cb720039754d738850c15ab5212a9e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HBox(children=(IntSlider(value=1000, description='Max Iters', max=3000, min=100, step=100), Che…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "widget = Toy2DWidget(\n",
    "    Toy2DCMP,\n",
    "    num_iters=1000,\n",
    "    epsilon=0.7,\n",
    "    problem_type=\"Convex\",\n",
    "    primal_optim=\"SGD\",\n",
    "    primal_lr=0.02,\n",
    "    dual_optim=\"SGD\",\n",
    "    dual_lr=0.5,\n",
    "    x=0.7,\n",
    "    y=1.8,\n",
    "    dual_restarts=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References <a class=\"anchor\" id=\"references\"></a>\n",
    "\n",
    "- J. Degrave and I. Korshunova. Why machine learning algorithms are hard to tune and how to fix it. Engraved,   [blog](www.engraved.blog/why-machine-learning-algorithms-are-hard-to-tune/), 2021. <a class=\"anchor\" id=\"blog1\"></a>\n",
    "- J. Degrave and I. Korshunova. How we can make machine learning algorithms tunable. Engraved,   [blog](https://www.engraved.blog/how-we-can-make-machine-learning-algorithms-tunable/), 2021. <a class=\"anchor\" id=\"blog2\"></a>\n",
    "- G. Gidel, H. Berard, G. Vignoud, P. Vincent, and S. Lacoste-Julien. A Variational Inequality\n",
    "Perspective on Generative Adversarial Networks. In ICLR, 2019.\n",
    "- G. M. Korpelevich. The extragradient method for finding saddle points and other problems. Matecon, 1976."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "e95f21d6d24b9c057dc89e3e94c7b1d285f49ff39a3b96b3de58d480bdf45a1c"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
